import numpy as np
import pandas as pd
from collections import defaultdict
import copy
import random
import matplotlib.pyplot as plt

random.seed()
np.random.seed()

# Load data from the CSV file using Pandas
data = pd.read_csv(r'C:\Users\grant\OneDrive - Oklahoma A and M System\School\PhD\Dissertation\P2 - Socail and Optimization\3 - Optimization\Test_Data\113,000\113k.csv')

# Extract data from csv into local variables
locations = np.array(data[['Latitude', 'Longitude']])
debris = data['Value'].tolist()
vulnerability = data['Priority'].tolist()

# Constants
NUM_SWARMS = 2  # Number of swarms
NUM_TRUCKS_PER_SWARM = 5  # Number of trucks per swarm
TRUCK_CAPACITY = 15  # truck (particle) capacity
VISIT_ALL_LOCATIONS = False

if VISIT_ALL_LOCATIONS:
    num_locations = len(locations)  # number of locations
    available_sites = list(range(len(locations)))
else:
    num_locations = len(np.where(np.array(debris) > 0)[0])  # number of locations
    available_sites = []

# Define PSO parameters
num_particles = NUM_SWARMS * NUM_TRUCKS_PER_SWARM
NUM_ITERATIONS = 100
INERTIA_WEIGHT = 0.8
C1 = 1.0  # personal best
C2 = 1.0  # global best

# Time calculation constants
TIME_PER_UNIT = 5  # 5 minutes per unit of debris
SPEED_MPH = 30  # Speed in miles per hour


def haversine_distance(lat1, lon1, lat2, lon2):
    # Convert latitude and longitude from degrees to radians
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])

    # Haversine formula
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = np.sin(dlat / 2.0) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2.0) ** 2
    c = 2 * np.arcsin(np.sqrt(a))

    # Radius of Earth in miles
    r = 3956
    return c * r


def calculate_time_for_one_site(start_loc, end_loc, debris_collected):
    distance = haversine_distance(start_loc[0], start_loc[1], end_loc[0], end_loc[1])
    time_for_trip = (distance / SPEED_MPH) * 60
    time_for_collection = debris_collected * TIME_PER_UNIT
    return time_for_trip + time_for_collection, distance


def update_velocity(current_velocity, personal_best_position, global_best_position, current_position):
    w = INERTIA_WEIGHT  # inertia weight
    r1, r2 = np.random.rand(), np.random.rand()
    new_velocity = (w * current_velocity +
                    C1 * r1 * (personal_best_position - current_position) +
                    C2 * r2 * (global_best_position - current_position))
    return new_velocity


def calculate_fitness(time_taken, total_debris_collected):
    return total_debris_collected / time_taken if time_taken > 0 else 0


def run_pso_iterations(swarm_starting_locations, available_sites_orig):
    best_iteration_distance = np.inf  # Keep track of the shortest distance
    iteration_costs = []  # List to store the cost of each iteration
    best_iteration_visits = None  # To store the best iteration visits for later use

    # Initialize velocities and personal bests
    velocities = {i: np.zeros(2) for i in range(num_particles)}
    personal_bests = {i: (np.inf, -1) for i in range(num_particles)}  # (cost, index)
    global_best = (np.inf, -1)  # (cost, index)

    # Keep track of visited sites across all swarms
    visited_and_removed = np.array(debris)  # Initialize with initial debris amounts

    # Main loop for PSO iterations
    for iteration in range(NUM_ITERATIONS):
        print("\n\n", "-" * 50)
        print("-" * 50)
        print(f"Start iteration {iteration}\n")

        iteration_visits = defaultdict(lambda: [(-2, 0, 0, 0)])
        available_sites = available_sites_orig.copy()
        debris_counter = copy.deepcopy(
            visited_and_removed)  # Create a copy of the debris to be modified in each iteration
        closed_site_counter = 0  # Number of sites where debris are cleared
        iteration_cost = 0  # cost (time) taken for this iteration

        while closed_site_counter < num_locations:
            for swarm in range(NUM_SWARMS):
                for truck in range(NUM_TRUCKS_PER_SWARM):
                    print(f"STARTING WITH TRUCK {truck} in swarm {swarm}")
                    truck_moves_count = 0
                    amount_taken_by_truck = 0
                    truck_id = swarm * NUM_TRUCKS_PER_SWARM + truck

                    while amount_taken_by_truck < TRUCK_CAPACITY and closed_site_counter < num_locations:
                        print("----")
                        print(
                            f"Truck {truck} before move {truck_moves_count + 1} contains {amount_taken_by_truck} of debris")

                        if not VISIT_ALL_LOCATIONS:
                            available_sites = np.where(debris_counter > 0)[0]

                        debris_dropoff_cost = 0  # Initialize with 0 to add only when truck returns to depot

                        if len(available_sites) > 0:
                            # Update velocity
                            current_position = locations[iteration_visits[truck_id][-1][0]] if \
                                iteration_visits[truck_id][-1][0] != -2 else swarm_starting_locations[swarm]
                            personal_best_position = current_position  # Simplified assumption
                            global_best_position = current_position  # Simplified assumption
                            velocities[truck_id] = update_velocity(velocities[truck_id], personal_best_position,
                                                                   global_best_position, current_position)

                            # Use velocity to determine next site (simplified example)
                            if np.linalg.norm(velocities[truck_id]) > 0:
                                direction = velocities[truck_id] / np.linalg.norm(velocities[truck_id])
                                potential_sites = locations + direction  # Move in the direction of velocity
                                distances = np.linalg.norm(potential_sites - current_position, axis = 1)
                                new_site = np.argmin(distances)
                            else:
                                new_site = np.random.choice(available_sites)

                            previous_site = iteration_visits[truck_id][-1][0]
                            start_loc = locations[previous_site] if previous_site != -2 else swarm_starting_locations[
                                swarm]
                            end_loc = locations[new_site] if new_site != -2 else swarm_starting_locations[swarm]

                            amount_available = min(TRUCK_CAPACITY, debris_counter[new_site])

                            if (amount_taken_by_truck + amount_available) > TRUCK_CAPACITY:
                                amount_to_take = TRUCK_CAPACITY - amount_taken_by_truck
                            else:
                                amount_to_take = amount_available

                            amount_taken_by_truck += amount_to_take
                            debris_counter[new_site] -= amount_to_take

                            if debris_counter[new_site] == 0:
                                closed_site_counter += 1

                            debris_pickup_cost, distance = calculate_time_for_one_site(start_loc, end_loc,
                                                                                       amount_to_take)
                            iteration_visits[truck_id].append((new_site, amount_to_take, debris_pickup_cost, distance))

                            if amount_taken_by_truck == TRUCK_CAPACITY or closed_site_counter == num_locations:
                                print("Heading back to starting location to dump all the debris")
                                debris_dropoff_cost, distance = calculate_time_for_one_site(end_loc,
                                                                                            swarm_starting_locations[
                                                                                                swarm],
                                                                                            amount_taken_by_truck)
                                iteration_visits[truck_id].append((-2, 0, debris_dropoff_cost, distance))

                            iteration_cost += debris_pickup_cost + debris_dropoff_cost

                            if debris_counter[new_site] == 0:
                                if VISIT_ALL_LOCATIONS:
                                    available_sites.remove(new_site)
                                print(f"Site {new_site} is closed.")
                                print(f"Closed site counter is {closed_site_counter}/{num_locations} sites.")
                                if closed_site_counter == num_locations:
                                    debris_dropoff_cost, distance = calculate_time_for_one_site(end_loc,
                                                                                                swarm_starting_locations[
                                                                                                    swarm],
                                                                                                amount_taken_by_truck)
                                    iteration_visits[truck_id].append((-2, 0, debris_dropoff_cost, distance))
                                    iteration_cost += debris_dropoff_cost

                    total_time_for_truck = sum([visit[2] for visit in iteration_visits[truck_id]])
                    total_debris_collected_by_truck = sum([visit[1] for visit in iteration_visits[truck_id]])
                    fitness = calculate_fitness(total_time_for_truck, total_debris_collected_by_truck)

                    if fitness > personal_bests[truck_id][0]:
                        personal_bests[truck_id] = (fitness, truck_id)
                    if fitness > global_best[0]:
                        global_best = (fitness, truck_id)

        total_distance_travelled = sum([visit[3] for truck_id in iteration_visits for visit in iteration_visits[truck_id]])
        if total_distance_travelled < best_iteration_distance:
            best_iteration_distance = total_distance_travelled
            best_iteration_visits = copy.deepcopy(iteration_visits)  # Save the best iteration visits

        iteration_costs.append(best_iteration_distance)

    # Plot the convergence rate
    plt.plot(range(NUM_ITERATIONS), iteration_costs)
    plt.xlabel('Iteration')
    plt.ylabel('Best Iteration Distance')
    plt.title('Convergence Rate of PSO Algorithm')
    plt.show()

    return best_iteration_visits, best_iteration_distance


# Initialize swarm starting locations
swarm_starting_locations = []
for swarm in range(NUM_SWARMS):
    starting_location_index = np.random.choice(num_locations)
    swarm_starting_locations.append(locations[starting_location_index])

# Get best iteration
best_iteration_visits, best_iteration_distance = run_pso_iterations(swarm_starting_locations, available_sites)

# Print starting locations of swarms in latitude and longitude
print("\nStarting Locations of Swarms:")
for swarm, starting_location in enumerate(swarm_starting_locations):
    print(f"Swarm {swarm + 1}: Latitude {starting_location[0]}, Longitude {starting_location[1]}")

# List sites visited by each truck in each swarm for the best iteration
print("\nSites Visited and Debris Collected by Each Truck for Best Iteration:")
total_debris_collected_by_swarms = [0] * NUM_SWARMS
total_time_taken_by_swarms = [0] * NUM_SWARMS
total_distance_travelled_by_swarms = [0] * NUM_SWARMS

for swarm in range(NUM_SWARMS):
    print(f"\nSwarm {swarm + 1}:")
    for truck in range(NUM_TRUCKS_PER_SWARM):
        truck_id = swarm * NUM_TRUCKS_PER_SWARM + truck
        truck_sites = best_iteration_visits[truck_id]
        truck_sites_without_return = np.array([int(i[0]) for i in truck_sites])
        print(f"\nTruck {truck_id + 1} visited the following sites: {truck_sites_without_return}")
        available_capacity = TRUCK_CAPACITY  # Initialize with total truck capacity
        total_time_for_truck = 0
        total_debris_picked_up_by_truck = 0
        total_distance_travelled_by_truck = 0

        print("Amounts collected and remaining capacity:")

        for site, debris_picked_up, pickup_cost, distance in truck_sites[1:]:
            if site == -2:
                print("Truck is returning to the depot to unload debris.")
                available_capacity = TRUCK_CAPACITY
            else:
                available_capacity -= debris_picked_up
                print(
                    f"  - Site {site}: Collected: {debris_picked_up} units, Remaining Truck Capacity: {available_capacity:.2f} units")
                total_debris_picked_up_by_truck += debris_picked_up

            total_time_for_truck += pickup_cost
            total_distance_travelled_by_truck += distance

        print(f"Total debris collected by Truck {truck_id + 1}: {total_debris_picked_up_by_truck:.2f} units")
        print(f"Total time for Truck {truck_id + 1}: {total_time_for_truck:.2f} minutes")
        print(f"Total distance travelled by Truck {truck_id + 1}: {total_distance_travelled_by_truck:.2f} miles")

        total_debris_collected_by_swarms[swarm] += total_debris_picked_up_by_truck
        total_time_taken_by_swarms[swarm] += total_time_for_truck
        total_distance_travelled_by_swarms[swarm] += total_distance_travelled_by_truck

# Print total amount collected and total time for each swarm
print("\nTotal Amount Collected and Total Distance for Each Swarm:")
for swarm in range(NUM_SWARMS):
    print(f"\nSwarm {swarm + 1}: Total Amount Collected = {total_debris_collected_by_swarms[swarm]} units")
    print(f"Swarm {swarm + 1}: Total Time = {total_time_taken_by_swarms[swarm]:.2f} minutes")
    print(f"Swarm {swarm + 1}: Total Distance = {total_distance_travelled_by_swarms[swarm]:.2f} miles")

# Calculate total for both swarms
total_debris_collected = sum(total_debris_collected_by_swarms)
total_time_taken = sum(total_time_taken_by_swarms)
total_distance_travelled = sum(total_distance_travelled_by_swarms)

# Print totals for both swarms
print("\nTotal for Both Swarms:")
print(f"Total Amount Collected = {total_debris_collected} units")
print(f"Total Time = {total_time_taken:.2f} minutes")
print(f"Total Distance = {total_distance_travelled:.2f} miles")
